from io import StringIO

import pandas as pd

from reportlab.lib import pagesizes, colors
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import SimpleDocTemplate, Paragraph, TableStyle, Table, PageBreak

from ai_quickstart.openai_translator.ai_translator.book import Book, ContentType
from ai_quickstart.openai_translator.ai_translator.utils import LOG


class Writer:
    def __init__(self):
        pass

    # 保存翻译内容
    def save_translated_book(self, book: Book, output_file_path: str = None, file_format: str = "PDF"):
        if file_format.lower() == 'pdf':
            self._save_translated_book_pdf(book, output_file_path)
        elif file_format.lower() == "markdown":
            self._save_translated_book_markdown(book, output_file_path)
        else:
            raise ValueError(f"传入文件格式错误")

    # 保存为PDF文件
    def _save_translated_book_pdf(self, book: Book, output_file_path: str = None):
        if output_file_path is None:
            output_file_path = book.pdf_file_path.replace(".pdf", f"_translated.pdf")

        LOG.info(f"pdf_file_path：{book.pdf_file_path}")
        LOG.info(f"开始翻译：{output_file_path}")

        # 注册字体文件
        font_path = "../fonts/simsun.ttc"
        pdfmetrics.registerFont(TTFont("SimSun", font_path))

        simsun_style = ParagraphStyle("SimSun", fontName="SimSun", fontSize=12, leading=14)

        # 创建一个PDF document对象
        doc = SimpleDocTemplate(output_file_path, pagesize=pagesizes.letter)
        styles = getSampleStyleSheet()
        story = []

        for page in book.pages:
            for content in page.contents:
                if content.status:
                    if content.content_type == ContentType.TEXT:
                        text = content.translation
                        para = Paragraph(text, simsun_style)
                        story.append(para)
                    elif content.content_type == ContentType.TABLE:
                        table = content.translation

                        # 以下因4o模型返回的格式问题，导致数据处理异常的修正方法
                        clean_data = "\n".join(line for line in table.split("\n") if not set(line) <= {"|", "-", " "})
                        df = pd.read_csv(StringIO(table), sep="|", skiprows=1, skipinitialspace=True)
                        # 清理 DataFrame 列名（去掉首尾空格）
                        df.columns = df.columns.str.strip()
                        # 删除无效列（Markdown 分隔线可能会导致多余的空列）
                        df = df.dropna(axis=1, how="all")  # 删除所有值均为空的列
                        table = df

                        table_style = TableStyle([
                            ('BACKGROUND', (0, 0), (-1, 0), colors.grey),  # 背景色
                            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),  # 文本颜色
                            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),  # 对齐方式
                            ('FONTNAME', (0, 0), (-1, 0), 'SimSun'),  # 表头字体为 "SimSun"
                            ('FONTSIZE', (0, 0), (-1, 0), 14),  # 表头字体大小14
                            ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
                            ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
                            ('FONTNAME', (0, 1), (-1, -1), 'SimSun'),  # 更改表格中的字体为 "SimSun"
                            ('GRID', (0, 0), (-1, -1), 1, colors.black)
                        ])

                        # table的内容
                        pdf_table = Table(table.values.tolist())
                        # table的样式
                        pdf_table.setStyle(table_style)
                        story.append(pdf_table)

            # 如果不是最后一页，则换页
            if page != book.pages[-1]:
                story.append(PageBreak())

        # 生成pdf文件
        doc.build(story)
        LOG.info(f"翻译完成：{output_file_path}")

    # 保存为Markdown文件
    def _save_translated_book_markdown(self, book: Book, output_file_path: str = None):
        if output_file_path is None:
            output_file_path = book.pdf_file_path.replace('.pdf', f"_translated.md")

        LOG.info(f"pdf_file_path：{book.pdf_file_path}")
        LOG.info(f"开始翻译：{output_file_path}")

        with open(output_file_path, 'w', encoding='utf-8') as output_file:
            for page in book.pages:
                for content in page.contents:
                    if content.status:
                        if content.content_type == ContentType.TEXT:
                            text = content.translation
                            output_file.write(text + '\n\n')
                        elif content.content_type == ContentType.TABLE:
                            table = content.translation
                            # header = '| ' + ' | '.join(str(column) for column in table.columns) + ' |' + '\n'
                            # separator = '| ' + ' | '.join(['---'] * len(table.columns)) + ' |\n'
                            # body = '\n'.join(['| ' + ' | '.join(str(cell) for cell in row) + ' |' for row in table.values.tolist()]) + '\n\n'
                            # output_file.write(header,separator,body)
                            output_file.write(table)

                if page != book.pages[-1]:
                    output_file.write('---\n\n')

        LOG.info(f"翻译完成：{output_file_path}")
